# python3.6

import random
import rtde_control
import rtde_receive
from paho.mqtt import client as mqtt_client
import numpy as np
import sys
import time
import threading
from ft300 import RobotiqFt300
import keyboard

# robot ip
ip = "192.168.3.140"     #真实机器人ip
# ip="192.168.237.129"     #ursim ip

#client parameters
broker = 'localhost'
port = 1883
topic = "ee_pose"
# generate client ID with pub prefix randomly
client_id = f'python-mqtt-{random.randint(0, 100)}'

# Parameters
velocity = 0.5
acceleration = 0.5
dt = 1.0/100  # 2ms
lookahead_time = 0.1
gain = 300

ratio_translation=0.5    #机器人末端平移与touch末端平移的比例
ratio_rotation=0.2       #机器人末端旋转与touch末端旋转的比例

# 设置计数器和更新频率
servo_count = 0
update_frequency = 100  # 每隔10次执行输出一次
start_time = 0
start_time = time.time()
distance=np.zeros(6)

pose=np.zeros(6)


def quaternion_to_rotation_vector(quaternion):
    """
    Convert quaternion to rotation vector.
    """
    q0, q1, q2, q3 = quaternion
    angle = 2 * np.arccos(q0)
    if angle < 1e-6:
        return np.zeros(3)
    axis = np.array([q1, q2, q3]) / np.sin(angle / 2)
    return angle * axis


def rotation_vector_to_quaternion(rotation_vector):
    """
    Convert rotation vector to quaternion.
    """
    angle = np.linalg.norm(rotation_vector)
    if angle < 1e-6:
        return np.array([1.0, 0.0, 0.0, 0.0])
    axis = rotation_vector / angle
    half_angle = angle / 2
    q0 = np.cos(half_angle)
    q1, q2, q3 = np.sin(half_angle) * axis
    return np.array([q0, q1, q2, q3])


def quaternion_multiply(q1, q2):
    """
    Multiply two quaternions.
    """
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2

    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2

    return np.array([w, x, y, z])

def on_key_event(e,ft:RobotiqFt300):
    if(e.name=='r'):
        ft.FtZero()
        print("FT sensor zeroed")


def connect_mqtt() -> mqtt_client:
    def on_connect(client, userdata, flags, rc):
        if rc == 0:
            print("Connected to MQTT Broker!")
        else:
            print("Failed to connect, return code %d\n", rc)

    client = mqtt_client.Client(client_id)
    client.on_connect = on_connect
    client.connect(broker, port)
    return client

#发送机器人末端姿态，用于计算touch力的方向
def publish(client: mqtt_client,ft:RobotiqFt300):
    global pose
    while True:
        msg=f"{pose[3]} {pose[4]} {pose[5]} {ft.forceTorque[0]} {ft.forceTorque[1]} {ft.forceTorque[2]} {ft.forceTorque[3]} {ft.forceTorque[4]} {ft.forceTorque[5]}"
        client.publish("robot_ee_pose&FT", msg)
        force=f"{ft.forceTorque[0]} {ft.forceTorque[1]} {ft.forceTorque[2]} {ft.forceTorque[3]} {ft.forceTorque[4]} {ft.forceTorque[5]}"
        client.publish("FT",force)
        time.sleep(0.001)


# def subscribe(client: mqtt_client):
#     def on_message(client, userdata, msg):
#         print(f"Received `{msg.payload.decode()}` from `{msg.topic}` topic")

#     client.subscribe(topic)
#     client.on_message = on_message

def subscribe(client: mqtt_client, rtde_c: rtde_control, rtde_r: rtde_receive):
    def on_message(client, userdata, msg):
        
        global servo_count
        global start_time
        global distance
        global pose

        message = np.array(
            list(map(float, msg.payload.decode().split())), dtype=float)
        delta_pose = message[0:6]
        button_state = message[6]
        
        if(button_state):
            delta_pose[0:3] = delta_pose[0:3] * 0.001*ratio_translation        #touch末端位移,单位毫米，*0.001化为米
            delta_pose[3:6]=delta_pose[3:6]*ratio_rotation                     #touch末端旋转向量
            #用四元数计算目标姿态
            quat_rotation=rotation_vector_to_quaternion(delta_pose[3:6])
            quat_pose=rotation_vector_to_quaternion(pose[3:6])
            # quat_target_pose=quaternion_multiply(quat_rotation, quat_pose) # 绕基座
            #quat_target_pose = quaternion_multiply(quat_pose, quat_rotation)    # 绕tcp末端
            quat_target_pose=quat_pose
            vector_target_pose=quaternion_to_rotation_vector(quat_target_pose)

            pose[0:3]+=delta_pose[0:3]
            pose[3:6]=vector_target_pose
            rtde_c.servoL(pose, velocity, acceleration,
                            dt, lookahead_time, gain)
           
        prefix="tcp pose:"
        # distance=distance+delta_pose
        sys.stdout.write("\r")
        formatted_pose = [f"{value:.5f}" for value in rtde_r.getActualTCPPose()]
        output_string = f"{prefix}{formatted_pose}"
        sys.stdout.write(output_string)
        # sys.stdout.write(str(rtde_r.getActualTCPPose()))
        # sys.stdout.write(str(distance))
        sys.stdout.flush()
        # print(rtde_r.getActualTCPPose())
        servo_count += 1
        

        # if servo_count % update_frequency == 0:
        #     elapsed_time = time.time() - start_time
        #     execution_rate = servo_count / elapsed_time
        #     print(
        #         f"ServoL has been executed {servo_count} times. Execution rate: {execution_rate:.2f} Hz.")
        #     start_time += elapsed_time
        #     servo_count = 0

    client.subscribe(topic)
    client.on_message = on_message


def run():

    global pose

    rtde_c = rtde_control.RTDEControlInterface(ip)
    rtde_r = rtde_receive.RTDEReceiveInterface(ip)
    rtde_c.moveJ([-1.57,-1.57,-1.57,-1.57,1.57,0],0.08,0.01)
    print(rtde_r.getActualTCPPose())
    pose=rtde_r.getActualTCPPose()
    client = connect_mqtt()
    ft=RobotiqFt300()
    keyboard.hook(on_key_event,ft)
    # 创建接收和发送消息的线程
    receive_thread = threading.Thread(target=subscribe, args=(client, rtde_c, rtde_r))
    send_thread = threading.Thread(target=publish, args=(client, ft))

    # 设置线程为守护线程，使得程序可以退出时结束这两个线程
    receive_thread.daemon = True
    send_thread.daemon = True
    
    try:
        # 启动两个线程
        receive_thread.start()
        send_thread.start()
        # subscribe(client, rtde_c, rtde_r)
        # publish(client,ft)
        client.loop_forever()
    except KeyboardInterrupt:
        print("interrupted by keyboard")
        rtde_c.servoStop()
        rtde_c.stopScript()
        rtde_c.disconnect()
        rtde_r.disconnect()
        ft.DisConnect()
        client.disconnect()
        sys.exit(0)
    


if __name__ == '__main__':
    run()

